# -*- coding: utf-8 -*-
"""
Created on Thu Nov 25 20:05:59 2021

@author: bwm
"""
import pandas as pd

class find:
    #获取总数据（看你们要什么数据可以再改）
    def __init__(self,data_all, date_all):
        
        self.data_all = data_all
        self.date_all = date_all

    def get_data_all(path):
        income_data = pd.read_excel(path,header=0)  ##读入文件
        time_list=income_data['时间'].unique() ##将时间这一列读成列表，删去重复值
        data_all={}
        for item in time_list:
            result=income_data[income_data["时间"].str.contains(item)]  #获取同一时间的所有公司的数据，形成一个新的dataframe
            del result["时间"] ##把时间这一列删除
            result=result.set_index("代码", inplace=False)  
            data_all[item]=result ##将时间作为key,result这个dataframe作为values写入字典中

        return data_all
    
    #获取交易日数据(报错函数)，直到下一个存在这个日期，end_date不在就减一
    def get_date_all(data_all,start_date,end_date):
        date = list(data_all.keys())
        date_all = date[date.index(start_date):date.index(end_date)+1]
        return date_all
    
    #数据获取函数（根据和日期和股票获得各类指标）,数据结构为字典套dataframe,获取数据函数为
    def get_data(self,date,stock,s_index):#数据获取函数（可能要用数据清洗）
        #输入一个data得到当天所有数据的dateframe，在dateframe中根据stock和s_index得到某个股票相应指标
    	return self.data_all[date].loc[stock][s_index]
    
    #获取周围几个交易日的日期
    def get_date(self,date_get,i):
        #例如get_date(now_date,-1)获取now_date减去一个交易日的日期（这里可能需要构建一个交易日期list，找到now_date的日期后索引减一得到前一个交易日的日期）
        return self.date_all[self.date_all.index(date_get)+i]
    
    #获取某一支股票在某个交易日之前的所有历史交易信息
    def get_hist_data(self,start_date,date, stock):
        all_date = find.get_date_all(self.data_all,start_date, date)
        hist_data = pd.DataFrame(columns = list(self.data_all['2020-01-02'].columns), index = [])
        idx = all_date.index(date) + 1
        for i in range(idx):
            # 给hist_data新增一行，这一行是该股票在某一天的交易信息
            hist_data.loc[all_date[i]] = list(self.data_all[all_date[i]].loc[stock])

        return hist_data